/// Test the load bandwidth in the global memory.

#include <cuda.h>
#include <stdio.h>
#include <stdlib.h>
#include <tuple>

#include "benchmark_config.h"
#include "cudaDeviceQuery.h"
#include "gpu_timer.h"

typedef uint64_t DataType;

constexpr int kIter = 512;
constexpr int kWarpPerBlock = 16;
constexpr int kMaxOffset = 0;
constexpr int kMaxStride = 32;

// Each block access separate subarray to avoid sharing in L2 Cache
__global__ void mem_bw(DataType *array, DataType *sink, int offset = 0,
                       int stride = 1) {
  // thread index
  uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
  DataType temp = 0;
  for (int iter = 0; iter < kIter; iter++) {
    temp += array[(iter * blockDim.x * gridDim.x + tid) * stride + offset];
  }
  // Trick compiler not to eliminate LDG, while avoiding uncessary writes.
  if (tid > blockDim.x * gridDim.x) {
    sink[tid] = temp;
  }
}

int main() {

  initializeDeviceProp(0);
  const int BlockNum = 2 * gDeviceParams.NumSM;
  const int ThreadsPerBlock = kWarpPerBlock * kWarpSize;
  uint64_t total_byte =
      BlockNum * ThreadsPerBlock * sizeof(DataType) * kIter * kMaxStride +
      kMaxOffset;
  
  printf("BlockNum, ThreadsPerBlock: %d, %d\n", BlockNum, ThreadsPerBlock);
  // memory footprint of the kernel
  float footprint_MByte = (float)BlockNum * ThreadsPerBlock * sizeof(DataType) *
                          kIter / 1024 / 1024;
  printf("Memory Footprint = %f (MB)\n", footprint_MByte);

  DataType *posArray_g;
  checkGPUErrors(cudaMalloc(&posArray_g, total_byte));
  checkGPUErrors(cudaMemset(posArray_g, 0, total_byte));
  DataType *sink_g;
  checkGPUErrors(cudaMalloc(&sink_g, kWarpSize * sizeof(DataType)));
  // Warpup
  for (int i = 0; i < WARMUPNUM; ++i) {
    mem_bw<<<BlockNum, ThreadsPerBlock>>>(posArray_g, sink_g);
  }

  int offset = 0;
  int stride = 1;
  GpuTimer timer;
  float total_time = 0;
  for (int i = 0; i < 20; ++i) {
    timer.start();
    mem_bw<<<BlockNum, ThreadsPerBlock>>>(posArray_g, sink_g, offset, stride);
    timer.stop_and_wait();
    total_time += timer.duration();
  }

  total_time /= 20.f;
  float bw_sec = footprint_MByte / total_time;
  float max_bw = (float)2 * gDeviceParams.MemFrequency / 1000 *
                 gDeviceParams.MemBitWidth / 8;
  printf("Average kernel BW: %f GB/s\n", bw_sec);

  checkGPUErrors(cudaFree(posArray_g));
  checkGPUErrors(cudaFree(sink_g));
  return 0;
}
